import test_config_list, subprocess, json, shutil, reentrancy, sys
from contracts import *
from values import MyGlobals

# 依赖检查
found_depend = False
try:
    import z3
except:
    print("\033[91m[-] python z3模块未安装")
    found_depend = True
if not (subprocess.call("solc --version", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0):
    print("\033[91m[-] python solc模块未安装")
    found_depend = True
if found_depend:
    sys.exit(1)

global debug, read_from_blockchain

# 检查版本
def checkVersion():
    # 获取contracts目录涉及到的所有solc版本
    values = test_config_list.contracts_version.values()
    # 使用set去重
    unique_values = set(values)
    # 转换为列表
    unique_values_list = list(unique_values)
    # 排序
    unique_values_list.sort()
    # 打印版本，经常卡住是在不行手动solc-select install xxx安装
    print("\n涉及到的solc版本有：", unique_values_list)

    # 检查solc版本
    print("\n", "==========开始检查并自动安装solc版本，如果卡死，请终止后重新执行==========")
    for value in unique_values_list:
        if not (subprocess.call(["solc-select", "use", value], shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0):
            if subprocess.call(["solc-select", "install", value], shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0:
                print("solc", value, "安装成功")
            else:
                print("solc", value, "安装失败")
        else:
            print("solc", value, "已存在")
    print("\n", "==========solc版本检查完成，开始测试合约==========")

# 批量测试合约
def testAll():
    # 统计漏洞合约数量
    count1 = 0
    # 统计正常合约数量
    count2 = 0
    # 循环测试合约
    for i, v in test_config_list.contracts_version.items():
        # print(i, v)
        # 调试信息
        print("[ ]\033[1m 开始测试合约 %s.sol " % i)
        # 合约名
        name = test_config_list.contracts_names[i]
        # sol源码路径
        SourcePath = "contracts/" + i + ".sol"
        # 构造编译后的路径
        FilePath = SourcePath.replace('/', '_').replace('.', '_')
        # 切换solc版本
        if (subprocess.call(["solc-select", "use", v], shell = False, stdout = subprocess.PIPE, stderr = subprocess.PIPE) == 0):
            print("[ ]\033[1m 切换solc版本为", v)
        else:
            print("[ ]\033[1m solc版本切换失败")
        # 编译合约
        compile_contract(SourcePath, FilePath)
        # 得到合约bytecode
        contract_code_path = FilePath + '/' + name + '.bin-runtime'
        # 如果没有得到bytecode
        if not os.path.isfile(contract_code_path):  
            print('\033[91m[-] 编译失败，合约 %s.sol 不存在或代码有错误\033[0m\n' % i )
            return
        # 得到合约sourcemap
        sourcemap_path = FilePath + '/' + "combined.json"
        # 如果sourcemap不存在
        if not os.path.isfile(sourcemap_path):  
            print('\033[91m[-] SourceMap 不存在\033[0m\n' )
            return
        with open(sourcemap_path, 'r') as fw:
            # 加载json
            injson = json.load(fw)
            # 拿到与本次检测有关的部分
            sourceMapJson = injson['contracts'][SourcePath + ':' + name]['srcmap-runtime']
            # 得到源代码中字节码的映射关系
            sourceMap = sourcemap_corresponding(sourceMapJson)
        # 读取合约源代码
        with open(contract_code_path,'r') as f: code = f.read(); f.close()
        # 处理源代码中的\n、\r、空格
        code = code.replace('\n','').replace('\r','').replace(' ','')
        # 开始检测
        result = reentrancy.check(code, SourcePath, 0, MyGlobals.read_from_blockchain, sourceMap, False)
        if result:
            print('\033[91m[ ]  合约 %s.sol 发现重入攻击漏洞! \033[0m\n' % i)
            count1 += 1
        else:
            print('\033[92m[ ] 合约 %s.sol 未发现重入攻击漏洞 \033[0m\n' % i)
            count2 += 1
        shutil.rmtree(FilePath)
        
    print('\033[92m[ ]  统计结果：漏洞合约 %d 个，正常合约 %d 个 \033[0m\n' % (count1, count2))

if __name__ == '__main__':
    # 检查版本，这玩意儿首次执行需要打开，全部安装完后就可以注释掉了
    # checkVersion()
    # 测试所有合约
    testAll()